# 统一导入工具包
import math
import torch
import os
import torch.nn as nn
import torch.nn.functional as F
import torchtext
import time
import numpy as np
import random
import matplotlib.pyplot as plt
from torchtext.data.utils import get_tokenizer
from torch.nn import TransformerEncoder, TransformerEncoderLayer
'''
torchtext.data.Field 声明处理数据的方式
参数说明：
    tokenize 分词处理
    init_token 定义开始字符
    eos_token 定义结束字符
    lower 小写化处理
'''
# 声明处理方式，主要包括分词和小写化处理
TEXT = torchtext.data.Field(tokenize=get_tokenizer("basic_english"),
                            init_token='<sos>',
                            eos_token='<eos>',
                            lower=True)
# 划分数据集
train_txt, val_txt, test_txt = torchtext.datasets.WikiText2.splits(TEXT, root='./datasets', train='wiki.train.tokens', validation='wiki.valid.tokens', test='wiki.test.tokens')

# 统计数据
print('train_txt tokens: %d' % len(train_txt.examples[0].text))
print('val_txt tokens:   %d' % len(val_txt.examples[0].text))
print('test_txt tokens:  %d' % len(test_txt.examples[0].text))
# 依据训练集构建词典
TEXT.build_vocab(train_txt)

# 查看词典
length = len(TEXT.vocab)
print('词表大小: %d' % length)
TEXT.vocab.stoi

# 获取当前设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def batchify(data, bsz):
    '''
    将数据划分为用于训练的批次
    '''
    # 将文本形式的数据用 token 的相应索引来表示
    data = TEXT.numericalize([data.examples[0].text])
    # 获取总的批次数目
    nbatch = data.size(0) // bsz
    # 去除剩余的部分，比如总长度为12，而批次大小为5，那么剩余的2个 token 将不会被包括在内
    data = data.narrow(0, 0, nbatch * bsz)
    # 根据批次大小，划分数据集
    data = data.view(bsz, -1).t().contiguous()
    return data.to(device)


batch_size = 20
eval_batch_size = 10
train_data = batchify(train_txt, batch_size)
val_data = batchify(val_txt, eval_batch_size)
test_data = batchify(test_txt, eval_batch_size)


bptt = 35 # 句子长度
def get_batch(source, i):
    '''
    把数据进一步切分成长度为35的序列，最后返回的 data:[35, batch_size] ,每一列表示一个连续的序列
    '''
    seq_len = min(bptt, len(source) - 1 - i)
    data = source[i:i+seq_len]
    target = source[i+1:i+1+seq_len].view(-1)
    return data, target

data, targets = get_batch(train_data, 0)
print(data.size())
print(targets.size()) # target 表示待预测的下一个正确的词，用于计算模型损失，进而更新参数

src = ''
for id in data[:, 0]:
    src = src + (' %s' % TEXT.vocab.itos[id])

src.strip()

tgt = ''
for id in targets[0::20]:
    tgt = tgt + (' %s' % TEXT.vocab.itos[id])

tgt.strip()


class PositionalEncoding(nn.Module):
    '''
    给原始序列添加位置编码
    '''

    def __init__(self, d_model, dropout=0.1, max_len=100):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        # 首先初始化为0
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        # sine 和 cosine 来生成位置信息
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        # 词经过嵌入层后，再加上位置信息
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

def subsequent_mask(sz):
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask
plt.figure(figsize=(5,5))
plt.imshow(subsequent_mask(20))
plt.show()

class TransformerModel(nn.Module):
    '''
    ntoken: 词表大小，用于构建嵌入层
    ninp: 模型维度
    nhead: 多头注意力机制中 head 数目
    nhid: 前馈神经网络的维度
    nlayers: TransformerEncoderLayer叠加层数
    '''

    def __init__(self, ntoken, ninp, nhead, nhid, nlayers, dropout=0.5):
        super(TransformerModel, self).__init__()
        self.model_type = 'Transformer'
        self.src_mask = None
        self.pos_encoder = PositionalEncoding(ninp, dropout)  # 位置编码
        encoder_layers = TransformerEncoderLayer(ninp, nhead, nhid, dropout)  # EncoderLayer
        self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers)  # Encoder
        self.encoder = nn.Embedding(ntoken, ninp)  # 嵌入层
        self.ninp = ninp  # 模型维度

        # decoder 用于将隐藏层的表示转化成词表中 token 的概率分布
        self.decoder = nn.Linear(ninp, ntoken)

    def forward(self, src):
        # 生成 mask ，保证模型只能看到当前位置之前的信息
        if self.src_mask is None or self.src_mask.size(0) != len(src):
            device = src.device
            mask = subsequent_mask(len(src)).to(device)
            self.src_mask = mask

        src = self.encoder(src) * math.sqrt(self.ninp)
        src = self.pos_encoder(src)  # 位置编码
        output = self.transformer_encoder(src, mask=self.src_mask, src_key_padding_mask=None)
        output = self.decoder(output)
        return output
criterion = nn.CrossEntropyLoss()

ntokens = len(TEXT.vocab.stoi) # 词表大小
emsize = 200 # 嵌入层维度
nhid = 200 # nn.TransformerEncoder 中前馈神经网络的维度
nlayers = 2 # 编码器中 nn.TransformerEncoderLayer 层数
nhead = 2 # 多头注意力机制中“头”的数目
dropout = 0.2 # dropout
model = TransformerModel(ntokens, emsize, nhead, nhid, nlayers, dropout).to(device)
# 学习率
lr = 2.0
# 随机梯度下降
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
# 动态调整学习率
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)
# xavier_normal_初始化参数
for p in model.parameters():
    if p.dim() > 1:
        nn.init.xavier_normal_(p)


def train():
    model.train()  # 训练模式，更新模型参数
    total_loss = 0.
    start_time = time.time()  # 用于记录模型的训练时长
    ntokens = len(TEXT.vocab.stoi)  # 词表大小
    for batch, i in enumerate(range(0, train_data.size(0) - 1, bptt)):
        # 获取批次数据
        data, targets = get_batch(train_data, i)
        optimizer.zero_grad()
        output = model(data)
        # 计算损失
        loss = criterion(output.view(-1, ntokens), targets)
        # 计算梯度
        loss.backward()
        # 梯度裁剪，防止梯度消失/爆炸
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
        # 优化参数
        optimizer.step()

        # 打印训练记录
        total_loss += loss.item()
        log_interval = 200
        if batch % log_interval == 0 and batch > 0:
            cur_loss = total_loss / log_interval
            elapsed = time.time() - start_time
            print('| epoch {:3d} | {:5d}/{:5d} batches | '
                  'lr {:02.2f} | ms/batch {:5.2f} | '
                  'loss {:5.2f} | ppl {:8.2f}'.format(
                epoch, batch, len(train_data) // bptt, scheduler.get_lr()[0],
                              elapsed * 1000 / log_interval,
                cur_loss, math.exp(cur_loss)))
            total_loss = 0
            start_time = time.time()

def evaluate(eval_model, data_source):
    eval_model.eval() # 评估模式，不更新模型参数，仅评估模型当前的表现
    total_loss = 0.
    ntokens = len(TEXT.vocab.stoi) # 词表大小
    with torch.no_grad():
        for i in range(0, data_source.size(0) - 1, bptt):
            data, targets = get_batch(data_source, i)
            output = eval_model(data)
            output_flat = output.view(-1, ntokens)
            total_loss += len(data) * criterion(output_flat, targets).item()
    return total_loss / (len(data_source) - 1)


best_val_loss = float("inf")
epochs = 3  # 共训练3个epoch
best_model = None

for epoch in range(1, epochs + 1):
    # for epoch in range(1, 2): # Kagging test
    epoch_start_time = time.time()
    # 训练过程
    train()
    # 验证过程
    val_loss = evaluate(model, val_data)
    # 打印验证结果
    print('-' * 89)
    print('| end of epoch {:3d} | time: {:5.2f}s | valid loss {:5.2f} | '
          'valid ppl {:8.2f}'.format(epoch, (time.time() - epoch_start_time),
                                     val_loss, math.exp(val_loss)))
    print('-' * 89)

    # 记录最佳模型
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        best_model = model

    # 调整学习率
    scheduler.step()

# # 保存模型
# if not os.path.exists('datasets/models'):
#     os.makedirs('datasets/models')
# torch.save({'state_dict': model.state_dict()}, 'datasets/models/best_model.pth.tar')

# 保存模型
if not os.path.exists('temp2/models'):
    os.makedirs('temp2/models')
torch.save({'state_dict': model.state_dict()}, 'temp2/models/best_model.pth.tar')

# 计算交叉熵损失
test_loss = evaluate(best_model, test_data)

# 计算困惑度
ppl = math.exp(test_loss)
print('=' * 40)
print('| End of training | test ppl {:8.2f}'.format(ppl))
print('=' * 40)

# 先实例化一个模型
model = TransformerModel(len(TEXT.vocab.stoi), ninp=200, nhead=2, nhid=200, nlayers=2, dropout=0.2).to(device)
# 模型加载训练好的参数
# checkpoint = torch.load('datasets/models/best_model.pth.tar')
checkpoint = torch.load('temp2/models/best_model.pth.tar')
model.load_state_dict(checkpoint['state_dict'])


history = 'it seems'

h = []
for w in history.split():
    h.append([TEXT.vocab.stoi[w]])

while (True):
    # 把列表转化成 tensor ，然后计算模型输出
    output = model(torch.tensor(h).to(device))
    # 获取概率最大的5个单词的 id
    idxs = output[-1].argsort(descending=True).view(-1)[:10]
    # 随机选择其中一个
    r = random.randint(0, 10)
    h.append([r])
    # 句子结束
    if TEXT.vocab.itos[r] == '.' or TEXT.vocab.itos[r] == '<eos>':
        break

# 将下标转化成句子
sent = ''
for w in h:
    sent += TEXT.vocab.itos[w[0]] + ' '

# out_path = './tmp/hypotheses.txt'
out_path = './temp2/hypotheses.txt'
# out_path = './submit/hypotheses.txt'
with open(out_path, 'w', encoding='utf-8') as f:
    f.write('history: ' + history + '\n')
    f.write('hypotheses: ' + sent + '\n')

print(sent)
plt.show()